/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/** \ingroup sfrhist \file

    The SED calculation routines of the sfrhist class.  (Historically,
    this was put in a separate files so it can be compiled with higher
    optimization, but there is no functionality left...) 
*/

// $Id$

#include "config.h"
#include "sfrhist.h"
#include "hpm.h"
#include "counter.h"
#include "random.h"
#include "functors.h"
#include "boost/shared_ptr.hpp"
#include "boost/function.hpp"
#include "boost/lambda/lambda.hpp"
#include "boost/lambda/construct.hpp"
#include "boost/lambda/casts.hpp"
#include "boost/lambda/bind.hpp"
#include "boost/lexical_cast.hpp"
#include "misc.h"
#include "snapshot.h"

using namespace  blitz;
using std::cout;
using std::cerr;
using std::endl;
using namespace boost;


/** Returns a functor which gives the age of the specified population
    in the specified object, according to preferences
    keywords. */
boost::shared_ptr<mcrx::functors::T_particle_functor>
mcrx::sfrhist::age_functor (int object, Snapshot::species s) const
{
  using namespace mcrx;
  using namespace boost;
  using namespace boost::lambda;

  functors::T_particle_functor f;

  // Two cases are different. Gas particles return a dummy functor
  // which is never used, nstar particles return the tform of the
  // particle itself so that it's never changed.
  if (s==Snapshot::gas)
    return boost::shared_ptr<functors::T_particle_functor>(new functors::T_particle_functor(functors::dummy_tform()));
  else if (s==Snapshot::nstar)
    f = functors::own_tform();
  else {

    // Now we go ahead and deal with the complicated cases: disk & bulge.

    // Remember that the functor should return *formation time*, that is
    // the negative of the age at the start of the simulation. Hence the
    // ubiquitous minus signs when generating the functors below.

    const string popname = Snapshot::species_name(s);
    // the keywords are numbered starting with 1
    const string objstr = lexical_cast<string>(object+1);
  

    // the star formation history of the particles is determined by
    // the keyword "xxxpopmode"
    const string mode = p.getValue(popname + "popmode"+objstr, string ());
    cout << "Population " << popname << " object "+objstr << " mode " << mode << endl;
    // additional keywords depend on what the mode actually is
    if (mode == "instantaneous") {
      //the keyword is the age of the population at the start of
      //the simulation
      const double age = p.getValue(popname + "popage"+objstr, double ());
      cout << "Population " << popname << " object "+objstr << " age " << age << endl;
      f= functors::constant(-age);
    }
    else if (mode == "absent") {
      // in this case the particles are dark, we define this as negative
      // age (not born yet... ;-)
      f= functors::constant(blitz::quiet_NaN(double()));
    }
    else if (mode == "constant") {
      // The age is the age of the oldest stars, and the pop is
      // assumed to have been forming stars continuously up until the
      // start of the simulation
      const double age = p.getValue(popname+"popage"+objstr, double ());
      cout << "Population " << popname << " object "+objstr << " max age " << age << endl;
      f= functors::stochastic_uniform(-age,0.);
    }
    else if (mode == "exponential") {
      try {
	// The age is the age of the oldest stars, and the SFR is
	// assumed to have been declining exponentially up until the
	// start of the simulation
	const double age = p.getValue(popname+"popage"+objstr, double ());
	const double tau = p.getValue(popname + "poptau"+objstr, double ());
	cout << "Population " << popname << " object "+objstr << " max age " << age << " tau " << tau << endl;
	f= functors::stochastic_exponential(-age, -tau);
      }
      catch (...) {
	cerr << "Error: Exponential SFR not specified correctly" << endl;
	exit (1);
      }	
    }
    else {
      cerr << "Error: Unknown " << popname << " popmode: "
	   << mode << " for object " << object << endl;
      exit (1);
    }
  }

  // Check for defined min/max AGE -- this is compared to the current
  // time, so here we need the snapshot age to convert ages to the
  // equivalent formation times
  const double min_age = p.defined("stellar_min_age")?
    p.getValue("stellar_min_age",double()):0;
  const double max_age = p.defined("stellar_max_age")?
    p.getValue("stellar_max_age",double()):blitz::huge(double());

  return boost::shared_ptr<functors::T_particle_functor>(new functors::T_particle_functor(functors::range_clip(f, snap.time()-max_age, snap.time()-min_age)));

}


/** Returns a functor which gives the metallicity of the specified
    population in the specified object, according to preferences
    keywords. The metallicity distribution is an exponential, without
    regard for population. */
boost::shared_ptr<mcrx::functors::T_particle_functor>
mcrx::sfrhist::z_functor (int object, Snapshot::species s) const
{
  const string popname = Snapshot::species_name(s);
  // the keywords are numbered starting with 1
  const string objstr = lexical_cast<string>(object+1);
  
  const double central_metallicity =
    p.defined("metallicity_gradient"+objstr)?
    p.getValue("central_metallicity" +objstr, double ()): 0;
  const double gradient= log(10.)* // gradient is given in 10, not e log
    (p.defined("metallicity_gradient"+objstr)?
     p.getValue("metallicity_gradient"+objstr, double ()): 0);

  cout << "Object " << objstr << " " << Snapshot::species_name(s)
       << " central metallicity " << central_metallicity
       << " gradient " << gradient << endl;

  return boost::shared_ptr<mcrx::functors::T_particle_functor>
    (new functors::T_particle_functor(functors::parent_radial_exponential(central_metallicity, gradient)));
}


vector<boost::shared_ptr<mcrx::sfrhist_snapshot::T_age_metallicity_functor> >
mcrx::sfrhist::compound_functors(Snapshot::species s) const
{
  vector<boost::shared_ptr<sfrhist_snapshot::T_age_metallicity_functor> > v;
  for (int i=0; i<nobject; ++i) 
    v.push_back(boost::shared_ptr<sfrhist_snapshot::T_age_metallicity_functor>(new sfrhist_snapshot::T_age_metallicity_functor(functors::compound(*age_functor(i, s), *z_functor(i, s)))));

  // if we've defined an extra set of keywords, those are for the free particles.
  if(p.defined("diskpopmode"+boost::lexical_cast<string>(nobject+1)))
    v.push_back(boost::shared_ptr<sfrhist_snapshot::T_age_metallicity_functor>(new sfrhist_snapshot::T_age_metallicity_functor(functors::compound(*age_functor(nobject, s), *z_functor(nobject, s)))));

  return v;
}
		

/** Assigns age and metallicity for the snapshot particles based on
    preferences keywords. */
void mcrx::sfrhist::assign_age_metallicity() 
{
  cout << "Calculating particle ages and metallicities" << endl;

  if(nobject==0) {
    if(snap.n(Snapshot::disk)+snap.n(Snapshot::bulge)==0)
      return;
    
    // see if metallicities and ages are already set in the particles
    bool age_set=true;
    bool z_set=false;
    const star_particle_set& disk
      (dynamic_cast<const star_particle_set&>(snap.set(Snapshot::disk)));
    for(int i=0; i<snap.n(Snapshot::disk); ++i) {
      age_set = age_set && (disk.tf(i)<0);
      z_set = z_set || (disk.z(i)>0);
    }
    const star_particle_set& bulge
      (dynamic_cast<const star_particle_set&>(snap.set(Snapshot::bulge)));
    for(int i=0; i<snap.n(Snapshot::bulge); ++i) {
      age_set = age_set && (bulge.tf(i)<0);
      z_set = z_set || (bulge.z(i)>0);
    }
    if(age_set && z_set) {
      cout << "Ages and metallicities of old star particles already set in snapshot\n";
      return;
    }

    if(!p.defined("diskpopmode1")) {
      cerr << "Error: Simulation has old star particles but no metallicity/age is defined.\n";
      exit(1);
    }
  }

  // The general idea here is to get the functors and then call the
  // assign_age_metallicity function of the snapshot object. The
  // different cases are handled slightly differently so it's a little
  // repetitive.
  snap.assign_age_metallicity(Snapshot::gas, 
			      compound_functors(Snapshot::gas));
  snap.assign_age_metallicity(Snapshot::disk, 
			      compound_functors(Snapshot::disk));
  snap.assign_age_metallicity(Snapshot::bulge, 
			      compound_functors(Snapshot::bulge));
  snap.assign_age_metallicity(Snapshot::nstar, 
			      compound_functors(Snapshot::nstar));

  if(p.defined("ascii_output_file")) {
    ofstream of(p.getValue("ascii_output_file",string()).c_str());
    snap.set(Snapshot::gas).write_ascii_table(of);
    snap.set(Snapshot::disk).write_ascii_table(of);
    snap.set(Snapshot::bulge).write_ascii_table(of);
    snap.set(Snapshot::nstar).write_ascii_table(of);
  }
};
